/*
 * SpeculationManager.cpp
 *
 *  Created on: Mar 28, 2012
 *      Author: user
 */

#include "SpeculationManager.h"
#include "MemoryManager.h"

#include <pthread.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <net/if.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdio.h>
#include <strings.h>
#include <string.h>
#include <time.h>
#include <errno.h>
#include "MemoryManager.h"

using namespace std;

static const int defaultTotalSpeculationSlots( 0 );
static const int defaultRunningSpeculationSlots( 0 );

static int g_port = 54756;

SpeculationManager* SpeculationManager::m_instance = 0;

SpeculationManager* SpeculationManager::Instance()
{
	if( m_instance == 0 )
	{
		m_instance = new SpeculationManager(defaultTotalSpeculationSlots, defaultRunningSpeculationSlots);
	}

	return m_instance;
}

SpeculationManager::SpeculationManager(int maxTasks, int maxRunningTasks)
	: m_maximumTasks( maxTasks ), m_maximumRunningTasks( maxRunningTasks ), m_rank( 0 )
{
    gettimeofday(&m_startTimeStamp,NULL);

    struct sockaddr_in serv_addr;

    m_listenSocket = socket(AF_INET, SOCK_STREAM, 0);

    int on = 1;
    setsockopt( m_listenSocket, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on) );

    memset( &serv_addr, 0, sizeof(serv_addr) );

    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = INADDR_ANY;
    serv_addr.sin_port = htons(g_port);

    if( bind( m_listenSocket, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0 )
    {
    	DEBUG_LOG( "Error opening master listen socket\n" )
        exit(0);
    }

    DEBUG_LOG("Master listen socket opened\n")

    listen(m_listenSocket, 5);

}

SpeculationManager::~SpeculationManager()
{
	DEBUG_LOG("*** SpeculationManager destructor ***\n")

	shutdown( m_listenSocket, SHUT_RDWR );
    close( m_listenSocket );
	killAllSpeculations();
}

void SpeculationManager::startSpeculation(SpeculationTask *speculation)
{
	speculation->setStatus( SPS_RUNNING );

	DEBUG_LOG( "Master starting speculation: %s\n", speculation->getSignature().c_str() )

	// The speculation is waiting to be executed by the manager
	// Kick it off and bail

	pid_t pID = fork();

	if( pID == 0 )
	{
        m_log = fopen("./specEx.log","a");

		// This is the clone

		// Open Socket

        DEBUG_LOG("Clone %d opening socket to master\n", getpid())

	    struct sockaddr_in serv_addr;
	    struct hostent *server;

	    speculation->m_socket = socket(AF_INET, SOCK_STREAM, 0);

        if( speculation->m_socket == -1 )
        {
        	DEBUG_LOG("Error clone %d socket: %d\n", getpid(), errno)
			exit(0);
        }

        DEBUG_LOG("Clone %d opened socket to master\n", getpid())

	    server = gethostbyname("127.0.0.1");

        DEBUG_LOG("Clone %d gethostbyname passed\n", getpid())

	    if (server == NULL)
	    {
	    	DEBUG_LOG("Error clone %d gethostbyname\n", getpid())
	        exit(0);
	    }

	    memset( &serv_addr, 0, sizeof(serv_addr) );

	    serv_addr.sin_family = AF_INET;

	    memcpy( &serv_addr.sin_addr.s_addr, server->h_addr, server->h_length );

	    serv_addr.sin_port = htons(g_port);

	    DEBUG_LOG("Clone %d connecting to %lx\n", getpid(), (unsigned long)serv_addr.sin_addr.s_addr)

	    while( connect( speculation->m_socket, (struct sockaddr *) &serv_addr, sizeof(serv_addr) ) < 0 )
	    {
	    	DEBUG_LOG("*** Error clone %d connecting to master: %d\n", getpid(), errno)
	    }

	    DEBUG_LOG("Clone %d connected\n", getpid())

		m_rank = 1;

		// Protect all the memory modifications
		protectAll();

		DEBUG_LOG("Clone %d is executing command: %s\n", getpid(), speculation->getSignature().c_str())

		// Launch the actual task
		speculation->execute();

		// I am a clone, signal to the master that the speculation has completed and wait here

		speculation->setStatus( SPS_COMPLETED );

		DEBUG_LOG("Clone %d finished executing the speculation: %s\n", getpid(), speculation->getSignature().c_str() )

		// The speculation has completed, signal the master
		int status = 1; // Speculation completed

		speculation->sendCommandToMaster( status );

		DEBUG_LOG("Clone %d has sent speculation complete message and is now waiting for ack\n", getpid() )

		commitSpeculation( speculation );

	}
	else
	{
		// This is the master

		// Wait for incoming connection on listening socket and open new socket

		DEBUG_LOG("Master waiting connection from clone %d\n", pID)

		socklen_t clilen;
		struct sockaddr_in cli_addr;

		clilen = sizeof(cli_addr);

		while(1) 
		{
			speculation->m_socket = accept( m_listenSocket, (struct sockaddr *) &cli_addr, &clilen );

			if( speculation->m_socket < 0 )
			{
				if( errno == EINTR )
				{
					continue;
				}
				else
				{
					DEBUG_LOG("*** Master accept() error connecting to clone %d: %d\n", pID, errno)
					exit( 0 );
				}
			}
			else
			{
				break;
			}
		}

		DEBUG_LOG("Master connected to clone\n")

		speculation->setProcessId( pID );

		DEBUG_LOG("Master started speculation clone: %d\n", pID)

	}

}

void SpeculationManager::commitSpeculation(SpeculationTask *speculation)
{
	if( m_rank == 0 )
	{
		// This is master

		int status;

		// If the speculation is still running, wait for it to complete
	    if( speculation->getStatus() == SPS_RUNNING )
	    {
			recvCommandFromClone( speculation, &status );

			DEBUG_LOG("Master received command\n")

			// The speculation was completed successfully
			if( status == 1 )
			{
				DEBUG_LOG("Master received completion command\n")
				speculation->setStatus( SPS_COMPLETED );
			}
	    }

	    DEBUG_LOG( "Master sending commit request to clone %d\n", speculation->getProcessId() )

		status = -1; // Send commit command
		sendCommandToClone( speculation, status );

		// Commit operations here

		int pageCount;
		recvCommandFromClone(speculation, &pageCount);

		DEBUG_LOG("Master receiving %d pages\n", pageCount)

		int pageSize = getPageSize();

		while (pageCount > 0)
		{
			void *pageAddress = 0;
			recvFromClone(speculation, sizeof(void*), &pageAddress);

//			DEBUG_LOG("Master received page with base address: %p\n", pageAddress)

			if (pageAddress == 0)
			{
				DEBUG_LOG("Received null page from clone, exiting\n")
				exit(1);
			}

			unsigned char* baseAddress;
			baseAddress = (unsigned char*)pageAddress;

			ssize_t ret = 0;
			int recvSize = 0;

			do
			{
				ret = read( speculation->m_socket, ((unsigned char*)baseAddress) + recvSize, pageSize - recvSize );
				if( ret != -1 )
				{
					recvSize += ret;
				}
			}
			while( (ret == -1 && errno == EINTR) || recvSize != pageSize);

			pageCount--;
		}

		// Wait for acknowledgement that commit has been performed
		DEBUG_LOG("Master all pages received, waiting for commit ack\n")

		recvCommandFromClone( speculation, &status );

		if( status == 2 )
		{
			DEBUG_LOG("Master received commit successful command\n")
		}
        else
        {
        	DEBUG_LOG("ERROR: Master received unknown command\n")
        }

		DEBUG_LOG("Master sending finish request to clone %d\n", speculation->getProcessId() )

		status = -2;
		sendCommandToClone( speculation, status );

		DEBUG_LOG( "Master removing clone %d\n", speculation->getProcessId() )

		waitpid( speculation->getProcessId(), &status, 0 );

		// Remove speculated task from list
		m_tasks.erase( speculation->getSignature() );
		delete speculation;
	}
	else
	{
		int status;

		// Wait for commit command from the master
		speculation->recvCommandFromMaster( &status );

		DEBUG_LOG("Clone %d received command from master\n", getpid() )

		// Unprotect the memory
		unprotectAll();

		if( status == -1 ) // Commit
		{
			DEBUG_LOG("Clone %d received commit request\n", getpid() )

			// Perform the commit here
			// TODO commit operations here

			const vector<void*>& dirtyPages = getDirtyPages();
			size_t size = dirtyPages.size();

			DEBUG_LOG("Clone sending %d pages\n", (int)size )

			// Send total page count to master
			speculation->sendCommandToMaster(size);
			int pageSize = getPageSize();

			// Send page addresses and content
			for (size_t i = 0; i < size; i++)
			{
				// Send page start address
				void *address = (void*)dirtyPages[i];

//				DEBUG_LOG("Clone sending page with address: %p\n", address )

				speculation->sendToMaster(sizeof(void*), &address);
				write(speculation->m_socket, dirtyPages[i], pageSize);
			}

			DEBUG_LOG("Clone %d sending commit finished\n", getpid() )

			status = 2; // Commit finished
			speculation->sendCommandToMaster( status );

			DEBUG_LOG("Clone %d waiting for finish command\n", getpid() )

			// Wait for end command from the master
			speculation->recvCommandFromMaster( &status );

		}

		DEBUG_LOG("Clone %d is terminating\n", getpid() )

		exit(0);
	}
}

SpeculationTask *SpeculationManager::getSpeculationTask(string signature)
{
	// Search for requested speculations
	map<string, SpeculationTask*>::iterator iter = m_tasks.find( signature );

	if( iter != m_tasks.end() )
	{
		DEBUG_LOG("getSpeculationTask found task: %s\n", (char*)iter->second->getSignature().c_str())

		SpeculationTask *speculation = iter->second;

		// Kill all other speculations
		killAllSpeculationsButOne( speculation );

		return speculation;
	}

	DEBUG_LOG("getSpeculationTask task not found: %s\n", signature.c_str())

	// Kill all speculations
	killAllSpeculations();

	return NULL;
}

void SpeculationManager::killAllSpeculations()
{
	DEBUG_LOG("Master is killing all speculation clones\n")

	map<string, SpeculationTask*>::iterator iter;

	for( iter = m_tasks.begin(); iter != m_tasks.end(); iter++ )
	{
		pid_t pid = iter->second->getProcessId();

		DEBUG_LOG("Master killing clone %d\n", pid)

		kill( pid, SIGKILL );
		delete iter->second;
	}

	m_tasks.clear();
}

void SpeculationManager::killAllSpeculationsButOne(SpeculationTask* s)
{
	DEBUG_LOG("Master is killing all but one speculation clones: %s\n", s->getSignature().c_str() )

	map<string, SpeculationTask*>::iterator iter;

	for( iter = m_tasks.begin(); iter != m_tasks.end(); iter++ )
	{
		if( iter->second != s )
		{
			pid_t pid = iter->second->getProcessId();

			DEBUG_LOG("Master killing clone %d\n", pid)

			kill( pid, SIGKILL );
			delete iter->second;
		}
	}

	m_tasks.clear();
	m_tasks.insert( pair<string,SpeculationTask*>(s->getSignature(), s) );
}

void SpeculationManager::run()
{
    unsigned int runningTasks = 0;

    // Test for speculation completion
    map<string, SpeculationTask*>::iterator iter;

    for( iter = m_tasks.begin(); iter !=  m_tasks.end(); iter++ )
    {
	    SpeculationTask *clone = iter->second;

	    if( clone->getStatus() == SPS_RUNNING )
	    {
		    runningTasks++;

		    fd_set rfds;
		    struct timeval tv;
		    int retval;

		    FD_ZERO( &rfds );
		    FD_SET( clone->m_socket, &rfds );

		    tv.tv_sec = 0;
		    tv.tv_usec = 0;

		    // Detect messages coming from clones
		    retval = select( clone->m_socket+1, &rfds, NULL, NULL, &tv );

		    // If a message has arrived a clone is likely ready to commit
		    if( retval > 0 )
		    {
			    int status = 0;

			    recvCommandFromClone( clone, &status );

			    DEBUG_LOG("Master received asynchronous command\n")

			    // The speculation was completed successfully
			    if( status == 1 )
			    {
			    	DEBUG_LOG("Master received asynchronous completion command\n")

				    clone->setStatus( SPS_COMPLETED );
				    runningTasks--;
			    }
			    else
			    {
			    	DEBUG_LOG("Master received unknown command! %d\n", status)
			    }
		    }
	    }
    }

    // If current amount of running speculations is less than
    // available resources generate a new speculation

    if( runningTasks < m_maximumRunningTasks && m_tasks.size() < m_maximumTasks )
    {
	    // Request a new speculation from the application
	    SpeculationTask *speculation = m_speculationGenerator();

	    // Check whether the same speculation is already running/waiting
	    map<string, SpeculationTask*>::iterator iter = m_tasks.find( speculation->getSignature() );

	    if( iter == m_tasks.end() )
	    {
	    	DEBUG_LOG("Master adding speculation %s\n", speculation->getSignature().c_str())

	    	// Add to running speculations
		    m_tasks.insert( pair<string, SpeculationTask*>( speculation->getSignature(), speculation ) );

		    // Run the new speculation
		    startSpeculation(speculation);
	    }
	    else
	    {
		    delete speculation;
	    }
    }
}

void SpeculationManager::sendCommandToClone( SpeculationTask *clone, int command )
{
	write( clone->m_socket, &command, sizeof(int) );
}

void SpeculationManager::recvCommandFromClone( SpeculationTask *clone, int *command )
{
	ssize_t ret = 0;
	int recvSize = 0;

	do
	{
		ret = read( clone->m_socket, ((unsigned char*)command) + recvSize, sizeof(int) - recvSize );
		if( ret != -1 )
		{
			recvSize += ret;
//			printf("--- Received %d of %d\n", recvSize, (int)sizeof(int));
		}
	}
	while( (ret == -1 && errno == EINTR) || recvSize != sizeof(int) );
}

void SpeculationManager::sendToClone( SpeculationTask *clone, int size, const void *data )
{
	write( clone->m_socket, data, size );
}

void SpeculationManager::recvFromClone( SpeculationTask *clone, int size, void *data )
{
	ssize_t ret = 0;
	int recvSize = 0;

	do
	{
		ret = read( clone->m_socket, ((unsigned char*)data) + recvSize, size - recvSize );
		if( ret != -1 )
		{
			recvSize += ret;
//			printf("--- Received %d of %d\n", recvSize, s);
		}
	}
	while ((ret == -1 && errno == EINTR) || recvSize != size);
}
